from pyqpanda import *
from pyqpanda.Visualization.circuit_draw import draw_qprog


def question1(input: str) -> list:
    """
    state为输入的二进制表示，state[0]为低位
    根据input设置state为0,1,2,3
    """
    state = [0, 0]
    if input == '00':
        state[1] = 0
        state[0] = 0
    elif input == '01':
        state[1] = 0
        state[0] = 1
    elif input == '10':
        state[1] = 1
        state[0] = 0
    elif input == '11':
        state[1] = 1
        state[0] = 1
    else:
        raise Exception('非法输入')

    qvm = CPUQVM()
    qvm.init_qvm()
    qubits = qvm.qAlloc_many(2)
    cbits = qvm.cAlloc_many(2)
    cbits[0].set_val(state[0])
    cbits[1].set_val(state[1])

    """
    构建选择分支
    """

    x_q0 = X(qubits[0])
    i_q0 = I(qubits[0])
    q0_controlled_x = QIfProg(cbits[0] == 1, x_q0, i_q0)

    x_q1 = X(qubits[1])
    i_q1 = I(qubits[1])
    q1_controlled_x = QIfProg(cbits[1] == 1, x_q1, i_q1)
    prog = QProg()
    prog << q0_controlled_x << q1_controlled_x << H(qubits[1]) << CNOT(qubits[1], qubits[0])

    # draw_qprog(prog, 'pic', filename='q1_'+ input+'.png')

    qvm.directly_run(prog)  # 执行量子程序
    """
        输出的量子比特排列为{q1,q0},即[00,01,10,11]
    """
    stat = qvm.get_qstate()
    return stat


"""
??? 以下不确定，所以约定如下
??? 第一个量子比特指qubits[0],第三个量子比特指qubits[2]
"""


def question2(angle: float) -> list:
    qvm = CPUQVM()
    qvm.init_qvm()
    qubits = qvm.qAlloc_many(3)
    cbits = qvm.cAlloc_many(2)

    """
        获得经典比特
        cbits[0] 为 qubits[1]的测量结果
        cbits[1] 为 qubits[2]的测量结果
    """
    prog = QProg()
    prog << RY(qubits[2], angle) << H(qubits[1]) << CNOT(qubits[1], qubits[0]) \
    << CNOT(qubits[2], qubits[1]) << H(qubits[2]) \
    << Measure(qubits[1], cbits[0]) << Measure(qubits[2], cbits[1])

    """
        根据经典比特将q0变成q3
    """
    controled_x_q0 = QIfProg(cbits[0] == 1, X(qubits[0]), I(qubits[0]))
    controled_z_q0 = QIfProg(cbits[1] == 1, Z(qubits[0]), I(qubits[0]))

    prog << controled_x_q0 << controled_z_q0

    # draw_qprog(prog, 'pic', filename='q2_' + str(angle) + '.png')

    """
    执行量子程序并获取状态
    """
    result = qvm.directly_run(prog)
    # print(result)
    stat = qvm.get_qstate()
    # print(stat)

    """
        计算第一个量子比特的振幅模方
    """
    c0 = result['c0']
    c1 = result['c1']
    index0 = c1 * 4 + c0 * 2
    index1 = index0 + 1
    a0 = stat[index0]
    a1 = stat[index1]
    return [abs(a0) ** 2, abs(a1) ** 2]


if __name__ == "__main__":
    # print(question2(0.0))
    # print(question2(0.5))
    # print(question2(1.0))
    # print(question2(1.5))
    # print(question2(2.0))
    # print(question2(2.5))

    # print(question1('00'))
    # print(question1('01'))
    # print(question1('10'))
    # print(question1('11'))
    pass
